package com.zhou.base.components.bs.common.util.id;

import cn.hutool.core.thread.ThreadUtil;
import com.baomidou.mybatisplus.core.toolkit.IdWorker;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.List;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

public class SnowflakeUtil {

    public static long getId() {
        return IdWorker.getId();
    }

    @Slf4j
    public static class RedisSnowflakeWorkerId {
        /**
         * key标识，可以随便定义
         */
        private final String idWorker = "WorkerId";
        /**
         * lua 脚本
         */
        private final String script = "redis.replicate_commands()\n" +
                "local now = redis.call('TIME')[1]\n" +
                "local idWordsKey = KEYS[1]\n" +
                "local serviceName = KEYS[2]\n" +
                "local sp = ':'\n" +
                "for i = 0, 1023 do\n" +
                "    local serviceKey = idWordsKey..sp..serviceName..sp..i\n" +
                "    if redis.call('SETNX', serviceKey, now) == 1 then\n" +
                "        redis.call('Expire', serviceKey, 60)\n" +
                "        return i;\n" +
                "    end\n" +
                "end\n" +
                "return -1";
        /**
         * 服务名称
         */

        @Value("${spring.application.name:idworker-serviceName}")
        private String serviceName;
        @Resource
        private RedisTemplate redisTemplate;

        /**
         * 获取机器标识号
         */
        public Long getWorkerIdNum() {
            // 实例化脚本对象
            DefaultRedisScript<Long> lua = new DefaultRedisScript<>();
            lua.setResultType(Long.class);
            lua.setScriptText(script);
            List<String> keys = Lists.newArrayList(idWorker, serviceName);

            // 获取序列号
            Long num = (Long) redisTemplate.execute(lua, keys, keys.size());
            String key = String.join(":", keys) + ":" + num;

            // -1 代表机器用完了，重试
            if (num < 0) {
                log.error("目前Id已用完，请重新启动试试");
                System.exit(0);
            }
            // 自动续期
            this.autoExpire(key);
            return num;
        }

        /**
         * 自动续期
         */
        private void autoExpire(String key) {

            ScheduledExecutorService executorService = new ScheduledThreadPoolExecutor(1,
                    ThreadUtil.newNamedThreadFactory("id_auto_expire-", true));
            log.info("自动续期id成功:{}", key);
            executorService.scheduleAtFixedRate(() -> {
                redisTemplate.expire(key, 60, TimeUnit.SECONDS);
            }, 0, 20, TimeUnit.SECONDS);
        }

        @PostConstruct
        public void initMybatisPlus() {
            if (serviceName == null || serviceName.length() == 0) {
                log.error("雪花算法初始化失败，【 app.idKey 】配置为空。");
                return;
            }
            long num = getWorkerIdNum();
            // 获取前5位
            long workerId = num >> 5;
            // 获取后5位
            long dataCenterId = num & (~(-1L << 5L));
            // 自定义初始化雪花算法
            IdWorker.initSequence(workerId, dataCenterId);
        }
    }
}
